"""
数据模型定义，使用Pydantic创建模型和验证
"""
from pydantic import BaseModel, Field
from typing import List, Dict, Any, Optional

# 数据库模式模型
class Column(BaseModel):
    name: str
    type: str
    description: str
    enum_values: Optional[Dict[str, str]] = None  # 添加枚举值映射

class Table(BaseModel):
    name: str
    columns: List[Column]

class DatabaseSchema(BaseModel):
    tables: List[Table]

# 聊天和查询模型
class ChatMessage(BaseModel):
    message: str
    options: Dict[str, Any] = {}

class MessageResponse(BaseModel):
    message_id: str
    response: str
    sql_generated: str
    results: List[Dict[str, Any]]

class ChatHistoryItem(BaseModel):
    message_id: str
    type: str
    content: str
    sql_generated: Optional[str] = None
    results: Optional[List[Dict[str, Any]]] = None
    timestamp: str

class ChatHistory(BaseModel):
    messages: List[ChatHistoryItem]

# 查询选项模型
class QueryOptions(BaseModel):
    allow_joins: bool = False
    complex_aggregations: bool = False
    max_results: int = 100

class OptionsUpdate(BaseModel):
    options: Dict[str, Any]
